package com.heiye.blog.ai.advisor.impl;

import com.heiye.blog.ai.domain.dos.ChatMessageDO;
import com.heiye.blog.ai.domain.mapper.ChatMessageMapper;
import com.heiye.blog.ai.model.dto.AIChatRequest;
import com.heiye.blog.ai.model.openrouter.OpenRouterAssistantMessage;
import com.heiye.blog.ai.model.vo.AiChatReqVO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.deepseek.DeepSeekAssistantMessage;
import org.springframework.transaction.support.TransactionTemplate;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author: heiye
 * @date: 2025/10/09 下午8:54
 * @version: v1.0.0
 * @description: 自定义打印流式日志 Advisor
 */
@Slf4j
public class CustomStreamLoggerAndMessage2DBAdvisor implements StreamAdvisor {

    private final ChatMessageMapper chatMessageMapper;
    private final AIChatRequest aiChatRequest;
    private final TransactionTemplate transactionTemplate;

    public CustomStreamLoggerAndMessage2DBAdvisor(ChatMessageMapper chatMessageMapper,
                                                  AIChatRequest aiChatRequest,
                                                  TransactionTemplate transactionTemplate) {
        this.chatMessageMapper = chatMessageMapper;
        this.aiChatRequest = aiChatRequest;
        this.transactionTemplate = transactionTemplate;
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
        // 对话 UUID
        String chatUuid = aiChatRequest.getChatId();
        // 用户消息
        String userMessage = aiChatRequest.getUserMessage();

        Flux<ChatClientResponse> chatClientResponseFlux = streamAdvisorChain.nextStream(chatClientRequest);
        // 创建 AI 流式回答聚合容器 (线程安全)
        AtomicReference<StringBuilder> fullContent = new AtomicReference<>(new StringBuilder());

        // 返回处理后的流
        return chatClientResponseFlux
                .doOnNext(response -> {
                    // 获取推理内容
                    String reasoningContent = null;

                    // 获取响应内容
                    AssistantMessage assistantMessage = response.chatResponse().getResult().getOutput();

                    // 判断推理内容类型
                    if (assistantMessage instanceof DeepSeekAssistantMessage deepSeekAssistantMessage) {
                        reasoningContent = deepSeekAssistantMessage.getReasoningContent();
                    } else if (assistantMessage instanceof OpenRouterAssistantMessage openRouterAssistantMessage) {
                        reasoningContent = openRouterAssistantMessage.getReasoningContent();
                    }

                    // 获取正式回答
                    String text = response.chatResponse().getResult().getOutput().getText();

                    // 逐块收集内容 若推理内容有值，则响应推理内容，否则，说明推理结束了，响应正式回答
                    String chunk;
                    if (Objects.nonNull(reasoningContent)) {
                        chunk = reasoningContent;
                    } else {
                        chunk = text;
                    }

                    log.info("## chunk: {}", chunk);

                    // 若 chunk 块不为空，则追加到 fullContent 中
                    if (chunk != null) {
                        fullContent.get().append(chunk);
                    }

                })
                .doOnComplete(() -> {
                    // 流式完成后打印完整回答
                    String completeResponse = fullContent.get().toString();
                    log.info("\n==== FULL AI RESPONSE ====\n{}\n========================", completeResponse);

                    // 开启编程式事务
                    transactionTemplate.execute(status -> {
                        try {

                            // 1. 存储用户消息
                            chatMessageMapper.insert(ChatMessageDO.builder()
                                    .chatUuid(chatUuid)
                                    .content(userMessage)
                                    // 用户消息
                                    .role(MessageType.USER.getValue())
                                    .createTime(LocalDateTime.now())
                                    .build());

                            // 2. 存储 AI 回答
                            chatMessageMapper.insert(ChatMessageDO.builder()
                                    .chatUuid(chatUuid)
                                    .content(completeResponse)
                                    // AI 回答
                                    .role(MessageType.ASSISTANT.getValue())
                                    .createTime(LocalDateTime.now())
                                    .build());

                            return true;
                        } catch (Exception ex) {
                            // 标记事务为回滚
                            status.setRollbackOnly();
                            log.error("", ex);
                        }
                        return false;
                    });
                })
                .doOnError(error -> {
                   // 出错时打印已收集的部分
                   String partialResponse = fullContent.get().toString();
                    log.error("## Stream 流出现错误，已收集回答如下: {}", partialResponse, error);
                });
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        // order 值越小，越先执行
        return 99;
    }
}
